#coding=utf-8

import os
import sys
import time
import hmac
import base64
import hashlib
import logging
import requests
import argparse
import urllib.parse
from ruamel.yaml import YAML


logging.basicConfig(
    level = logging.INFO,
    format="[%(asctime)s] %(name)s : %(levelname)s : %(message)s",
    datefmt="%Y.%m.%d %H:%M:%S",
)


class MSG:
    """ 
    钉钉群机器人消息通知类
    逻辑：
        1. 利用钉钉群机器人，创建可主动发送消息 webhook 地址
        2. 利用 markdown 发送格式化消息。

    功能
        1. 便捷发通知。（一次运行即可发一次通知）
        2. 可以配置多个机器人（钉钉群），
        3. 添加机器人
        4. 查看现有的机器人
        5. 授权用户，防滥用
        6. 消息支持 markdown 格式。（包括链接，图片）
        7. 支持一些特殊表情，如: 🎈

    配置项：
        1. 机器人列表 {"机器人名称": "access_token 地址"}
        2. 消息标题 self.title
        3. 用户（发送消息者标识）

    注意：
        access_token 来自 webhook 地址
    """
    def __init__(
        self,
        user: str = "kb",  # 后续采用强制的初始化，跟踪调用消息发布者
        robot_name: str = "kb个人机器人",
        config: str = os.path.join(os.path.dirname(sys.argv[0]), 'config.yaml')
    ):
        """ 初始化配置项
            user:    消息发送的操作者
            robot_name:   选择发消息的机器人。（配置文件内的机器人名称）
            config:  机器人列表文件， 机器人名称及其对应的 access_token。
                    （注意，此处的机器人名称可以与群里机器人名称不一致）
        """
        logging.debug(f'初始化 {self.__class__}')
        logging.debug(f'初始化参数 user : {user}')
        logging.debug(f'初始化参数 robot_name : {robot_name}')
        logging.debug(f'初始化参数 config : {config}')
        self.User = user
        self.Headers = {
            'Content-Type': 'application/json; charset=utf-8'
        }
        self.Main_webhook = 'https://oapi.dingtalk.com/robot/send?access_token='
        self.Title = '🎈消息提醒'
        self.Json = {
            "msgtype": "markdown",
            "markdown": {
                "title": self.Title, 
                "text": ""
            }
        }
        self.Config = config
        self.Robot_name = robot_name
        self.Cfg_load = {}
        self.Robots = {}
        self.Authorized = {}
        self.Yaml = YAML()
        self.load_cfg(config)

    def load_cfg(self, cfg: str):
        """ 加载 yaml 类型配置文件, 设置 access_token 地址 """
        logging.debug(f'加载配置文件: {cfg}')
        if os.path.isfile(cfg):
            logging.debug('文件大小', os.stat(cfg).st_size)
        if os.path.isfile(cfg) and os.stat(cfg).st_size:
            with open(cfg, encoding='utf8') as C:
                self.Cfg_load = self.Yaml.load(C)
            if 'robots' in self.Cfg_load:
                self.Robots = self.Cfg_load['robots']
            if 'authorized' in self.Cfg_load:
                self.Authorized = self.Cfg_load['authorized']
        else:
            logging.error(f"文件不存在: {cfg}")
            print(f"\n{'='*80}\n{'='*80}\n===                ", end='')
            print("如果是首次使用，请先添加机器人和授权用户", end='')
            print(f"                  ===\n{'='*80}\n{'='*80}\n")

    def format_msg_json(self, msg: str, title: str):
        """ 将消息放入 json 内，并在消息内容前方插入 markdown 标题 """
        logging.debug('整理即将发送的消息')
        # 输入的 title 优先使用
        if title.strip() != "":
            self.Json['markdown']['title'] = title
        elif self.Title != self.Json['markdown']['title']:
            self.Json['markdown']['title'] = self.Title
        # 处理 msg 中的异常的转义字符
        msg = msg.replace(r'\n', '\n')
        msg = msg.replace(r'\t', '\t')
        self.Json['markdown']['text'] =  f"### **{self.Title}**\n{msg}"
        logging.info(f"最终消息文本: \n{self.Json['markdown']['text']}\n")

    def get_sign(self, secret: str):
        """ 获得 sgin 验证 """
        timestamp = str(round(time.time() * 1000))
        secret_enc = secret.encode('utf-8')
        string_to_sign = f'{timestamp}\n{secret}'
        string_to_sign_enc = string_to_sign.encode('utf-8')
        hmac_code = hmac.new(secret_enc, string_to_sign_enc, digestmod=hashlib.sha256).digest()
        return timestamp, urllib.parse.quote_plus(base64.b64encode(hmac_code))

    def get_webhook(self) -> str:
        """ 获得完整的 webhook 地址 """
        logging.debug('获得完整的 webhook 地址')
        if self.Robot_name in self.Robots:
            access_token = self.Robots[self.Robot_name]['access_token']
            if access_token.isalnum():
                if self.Robots[self.Robot_name]['secret']:
                    timestamp, sign = self.get_sign(self.Robots[self.Robot_name]['secret'])
                    return f"{self.Main_webhook}{access_token}&timestamp={timestamp}&sign={sign}"
                return f"{self.Main_webhook}{access_token}"
            else:
                logging.error(f'access_token 为空或存在特殊字符: {access_token}')
                sys.exit(1)
        else:
            logging.error(f'机器人名称未配置: {self.Robot_name}')
            sys.exit(1)
    
    def robot_check(self):
        """ 确认机器人已配置 """
        logging.debug('机器人确认')
        if self.Robot_name not in self.Robots:
            logging.error(f'机器人名称未配置: {self.Robot_name}')
            sys.exit(1)

    def authorized_check(self):
        """ 授权确认 """
        logging.debug('授权确认')
        if self.User not in self.Authorized:
            logging.error(f'未经授权的用户: {self.User}')
            sys.exit(1)

    def send(self, msg: str, title: str=''):
        """ 发送 msg 参数的文本消息，返回消息发送状态 """
        # 消息内容为空字符串时，不发送消息
        if msg.strip() == "":
            logging.warning("空消息不发送")
            return "空消息不发送"
        self.authorized_check()
        self.robot_check()
        self.format_msg_json(msg, title)
        webhook = self.get_webhook()
        logging.debug('发送消息')
        res = requests.post(webhook, headers=self.Headers, json=self.Json).json()
        if res['errcode'] == 0:
            logging.info(f"{self.User} 消息发送成功")
        else:
            logging.error(f"{self.User} 消息异常 {res['errmsg']}")

    def add_robot(self, name: str, access_token: str, secret: str='', save: bool=True) -> bool:
        """ 
        添加一个钉钉机器人 （成功时会写入配置文件，后续可以根据名字直接使用）
            name: 机器人名称。（自定义名称, 但不能与已存在的名称冲突）
            access_token:  机器人的 access_token 地址.
            secret:  机器人的 secret 加签密钥.
        """
        logging.debug(f'添加机器人 {name}: {access_token}')
        if 'robots' not in self.Cfg_load:
            self.Cfg_load['robots'] = {}
        if name in self.Cfg_load['robots']:
            logging.error(f'机器人名称已存在 {name}')
        elif access_token.isalnum():
            self.Cfg_load['robots'][name] = {
                'access_token': access_token,
                'secret': secret
            }
            if save:
                logging.warning(f"成功添加机器人 {name} : {access_token}")
                self.cfg_update()
            return True
        else:
            logging.error(f'access_token 为空或存在特殊字符:{access_token}')
        return False

    def add_robots(self, rbs: dict):
        """ 
        批量添加机器人（成功时会写入配置文件）
            rbt 三元数据，secret 没有可为空。示例如下:
            rbt = [
                [robot1名字, access_token, secret],
                [robot2名字, access_token, secret],
            ]
        """
        fail_rbt =[]
        for name, access_token, secret in rbs:
            if not self.add_robot(name, access_token, secret, False):
                fail_rbt.append(name)
        if fail_rbt:
            logging.warning(f"共 {len(rbs)} 个机器人，{len(fail_rbt)} 个机器人添加异常。异常机器人列表:\n\t{fail_rbt}")
        self.cfg_update()

    def add_user(self, username: str, description: str):
        """ 
        添加授权用户 
            username: 用户名称（发消息需要验证名称）
            description: 用户描述信息
        """
        add_ok = False
        if not username.strip() or not description.strip():
            logging.error("用户名或描述为空！")
        elif 'authorized' in self.Cfg_load:
            if username in self.Cfg_load['authorized']:
                logging.warning(f"{username} 用户已存在")
            else:
                self.Cfg_load['authorized'][username] = description
                add_ok = True
        else:
            self.Cfg_load['authorized'] = {username: description}
            add_ok = True
        if add_ok:
            logging.warning(f'成功添加用户 {username} : {description}')
            self.cfg_update()

    def del_robot(self, robot_name: str):
        """ 删除机器人 """
        if 'robots' not in self.Cfg_load or not self.Cfg_load['robots']:
            print('\n没有机器人可以删除\n')
            return
        if robot_name in self.Cfg_load['robots']:
            del self.Cfg_load['robots'][robot_name]
            logging.warning(f'成功删除机器人: {robot_name}')
            self.cfg_update()

    def robot_list(self):
        """ 列出已有的机器人 """
        logging.debug('列出现有的机器人')
        if 'robots' not in self.Cfg_load:
            print('\n\t还没有配置机器人\n')
            return
        all_robot_list = self.Cfg_load['robots'].keys()
        list_str = "\n".join([f'<{r}>' for r in all_robot_list])
        print(f"\n现有 {len(all_robot_list)} 个机器人:\n{list_str}")

    def cfg_update(self):
        """ 更新配置文件 """
        logging.debug(self.Cfg_load)
        logging.warning('更新配置文件')
        with open(self.Config, 'w', encoding='utf-8') as C:
            self.Yaml.dump(self.Cfg_load, C)


if __name__ == '__main__':
    arg = argparse.ArgumentParser()
    arg.add_argument('-u', '--user', 
        default='kb', help='授权用户名称'
    ) 
    arg.add_argument('-r', '--robot_name', 
        default='kb个人机器人', help='机器人名字'
    ) 
    arg.add_argument('-c', '--config', 
        default='config.yaml', help='配置文件'
    )
    arg.add_argument('-m', '--message', 
        default='',
        type=str,
        help='需要发送的消息内容 (支持部分 markdown 语法)',
    )
    arg.add_argument('--add_robot', 
        default='', 
        help='需要添加的机器人, 格式: robot_name:access_token'
    )
    arg.add_argument('--robot_secret', 
        default='',
        help='机器人安全设置的加签密钥'
    )
    arg.add_argument('--add_user', 
        default='', 
        help='添加授权用户, 格式: user_name:description'
    )
    arg.add_argument('--robot_list', 
        default=False,
        action='store_true', 
        help='列出现有的机器人'
    )
    opts = arg.parse_args()
    M = MSG(opts.user, opts.robot_name, opts.config)
    # 列出现有的机器人
    if opts.robot_list:
        M.robot_list()
    # 添加机器人
    if opts.add_robot:
        if opts.add_robot.count(':') == 1:
            M.add_robot(*opts.add_robot.split(':'), secret=opts.robot_secret)
        else:
            logging.error(f'--add_robot 格式错误: {opts.add_robot}')
    # 添加用户
    if opts.add_user:
        if opts.add_user.count(':') == 1:
            M.add_user(*opts.add_user.split(':'))
        else:
            logging.error(f'--add_user 格式错误: {opts.add_user}')
    # 发送消息
    if opts.message.strip():
        M.send(opts.message)
